import os
import sqlite3
import shutil
import zipfile
from pathlib import Path
import sys
import argparse

def check_and_connect_db(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 设置返回字典格式的结果
    cursor = conn.cursor()
    return conn, cursor

def unzip_file(zip_path, extract_to):
    """解压ZIP文件"""
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)

def process_aim_zip(aim_zip, db_directory, scan_directory, db_conn, cursor):
    print(f"开始处理: {aim_zip}")
    
    # 根据aim_zip在db_directory下创建相应的文件夹
    folder_path = (Path(db_directory) / aim_zip).with_suffix('')

    if not folder_path.exists():
        folder_path.mkdir(parents=True, exist_ok=True)
        print(f"已创建文件夹: {folder_path}")
    else:
        print(f"文件夹已存在: {folder_path}")

    # 如果存在与aim_zip同名的ZIP文件（位于db_directory），则先解压缩到对应文件夹
    zip_path = Path(db_directory) / aim_zip
    if zip_path.exists() and zipfile.is_zipfile(zip_path):
        unzip_file(str(zip_path), str(folder_path))
        print(f"已解压ZIP文件: {zip_path} 到 {folder_path}")

    # 更新所有对应此aim_zip的记录，将deal_flag设置为1
    cursor.execute("UPDATE photo_move SET deal_flag=1 WHERE aim_zip=?", (aim_zip,))
    db_conn.commit()

def main(db_directory, scan_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    # 查询photo_move表中所有不同的aim_zip值
    cursor.execute("SELECT DISTINCT aim_zip FROM photo_move WHERE deal_flag=0")
    aim_zips = cursor.fetchall()

    if not aim_zips:
        print("没有需要处理的文件夹")
        return

    for aim_zip_row in aim_zips:
        aim_zip = aim_zip_row['aim_zip']
        process_aim_zip(aim_zip, db_directory, scan_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process files based on photo_move table and manage folders.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")
    parser.add_argument('--scan-dir', required=True, help="Directory to scan for files")

    args = parser.parse_args()

    db_directory = args.db
    scan_directory = args.scan_dir

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    if not os.path.isdir(scan_directory):
        print(f"错误: {scan_directory} 不是一个有效的目录")
        sys.exit(1)

    main(db_directory, scan_directory)